import cv2
from ultralytics import YOLO
import torch

class YoloPredictor:
    """
    一个YOLO预测器类，用于加载模型、执行预测、跟踪和实例分割。
    """
    def __init__(self, model_path):
        """
        初始化预测器，加载YOLO模型。

        :param model_path: YOLO模型文件的路径（.pt文件）。
        """
        # 检查是否有可用的GPU
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        print(f"正在使用设备: {self.device}")
        self.model = YOLO(model_path,task='detect')
        self.results = None

    def predict_and_track(self, frame):
        """
        对单帧图像执行目标检测、实例分割和跟踪。

        :param frame: 从摄像头捕获的图像帧。
        :return: Ultralytics的Results对象。
        """
        # 使用 model.track() 同时进行检测、分割和跟踪
        # persist=True 告诉跟踪器当前图像是序列中的下一帧
        # verbose=False 可以减少不必要的终端输出
        
        # self.results = self.model.track(frame, persist=True, verbose=False)
        self.results = self.model.predict(frame, verbose=False,show=True)
        return self.results

    def visualize_results(self):
        pass
    def print_results(self):
        """
        将当前帧的检测、分割和跟踪结果打印到终端。
        """
        if not self.results or not self.results[0]:
            return

        result = self.results[0]
        boxes = result.boxes
        masks = result.masks
        class_names = result.names

        print("--- 实时检测结果 ---")
        if boxes is not None and len(boxes) > 0:
            for i in range(len(boxes)):
                class_id = int(boxes.cls[i].item())
                class_name = class_names.get(class_id, "未知类别")
                confidence = boxes.conf[i].item()
                
                track_id_info = ""
                # 检查是否有跟踪ID
                if boxes.id is not None:
                    track_id = int(boxes.id[i].item())
                    track_id_info = f", 跟踪ID: {track_id}"

                print(f"  物体 {i+1}: 类别='{class_name}', 置信度={confidence:.2f}{track_id_info}")

            if masks is not None and masks.data.numel() > 0:
                print(f"  -> 检测到 {len(masks)} 个物体的实例分割掩码。")
            else:
                print("  -> 未进行实例分割（或模型不支持）。")
        else:
            print("  当前帧未检测到任何物体。")
        
        print("-" * 22 + "\n")


def main():
    # 打开默认摄像头
    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        print("错误：无法打开摄像头。")
        return

    print("\n摄像头已启动，正在进行实时检测... 按 Ctrl+C 退出。")
    """
    主函数，用于启动摄像头并进行实时预测。
    """
    # --- 配置 ---
    # 请确保您的自定义模型路径正确，这里假设权重文件是 best.pt
    # 如果您的模型文件不是 best.pt，请修改下面的路径
    # 注意：请将此路径更改为您计算机上的实际路径
    model_path = r"weights/yolo11x-200-epochs/best.pt"
    
    # 初始化YOLO预测器
    try:
        predictor = YoloPredictor(model_path)
    except Exception as e:
        print(f"错误：无法加载模型，请检查路径 '{model_path}' 是否正确。错误信息: {e}")
        return
    try:
        while True:
            # 读取一帧
            ret, frame = cap.read()
            if not ret:
                print("错误：无法从摄像头捕获帧。")
                break
            hsv_image = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
            h, s, v = cv2.split(hsv_image)
            s = cv2.add(s, 50)
            final_hsv = cv2.merge([h, s, v])
            final_bgr = cv2.cvtColor(final_hsv, cv2.COLOR_HSV2BGR)
            # 执行预测和跟踪
            predictor.predict_and_track(final_bgr)

            # 在终端中打印结果
            predictor.print_results()

    except KeyboardInterrupt:
        print("\n检测已停止。")
    finally:
        # 释放摄像头并关闭所有窗口
        cap.release()
        cv2.destroyAllWindows()
        print("资源已释放。")

if __name__ == "__main__":
    main() 